import datetime
import os
import time
import ujson
import logging
from absl import app, flags
import traceback
from typing import List, Dict, Optional, Set
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import coin.proto.coin_request_pb2 as cr_pb2
import coin.proto.coin_meta_service_pb2 as cms_pb2
import coin2.service.account.account_info_service_pb2 as ais_pb2
import coin2.service.strategy.trading_summary_service_pb2 as tss_pb2
import coin2.service.feed.interval_feed_service_pb2 as fs_pb2
import coin2.service.web3.order_service_pb2 as web3_pb2
import xunkemgmt.service.market_quotes_service_pb2 as mqs_pb2
from xunkemgmt_client.base.datetime_util import (to_timestamp_int, to_datetime)
from xunkemgmt_client.client.api_client import XunkemgmtClient
from python.experimental.lucasyoo.accounting.offset_accounting2 import (Data)
import python.experimental.lucasyoo.accounting.accounting as lac

matplotlib.use("Agg")
os.environ["TZ"] = "GMT"
time.tzset()

FLAGS = flags.FLAGS

REAL_TRANSFER_START_DT = datetime.datetime(2023, 1, 1)  # Start historical Transfer

MASTER_ACCOUNTING_ID = {
    "Binance": {
        "account_id": 239,
        "owner": "prestoinv0",
        "market_type": "Spot",
        "quote": "USDT"
    },
    "Okex": {
        "account_id": 526,
        "owner": "prestoinv0",
        "market_type": "ExWallet",
        "quote": "USDT"
    },
    "Huobi": {
        "account_id": 296,
        "owner": "prestoinv0",
        "market_type": "Spot",
        "quote": "USDT"
    },
    "Mexc": {
        "account_id": 1429,
        "owner": "prestoinv0",
        "market_type": "Spot",
        "quote": "USDT"
    },
    "Gateio": {
        "account_id": 1330,
        "owner": "prestoinv0",
        "market_type": "Spot",
        "quote": "USDT"
    },
    "Phemex": {
        "account_id": 1571,
        "owner": "prestoinv0",
        "market_type": "Spot",
        "quote": "USDT"
    },
    "Kucoin": {
        "account_id": 1430,
        "owner": "prestoinv0",
        "market_type": "ExWallet",
        "quote": "USDT"
    },
    "Bybit": {
        "account_id": 1402,
        "owner": "prestoinv0",
        "market_type": "Spot",
        "quote": "USDT"
    },
}


class Discretionary(object):
  def __init__(self, start_time: datetime.datetime, end_time: datetime.datetime,
               client: XunkemgmtClient, universe: List[str], output_dir: str):
    self._start_time: datetime.datetime = start_time
    self._end_time: datetime.datetime = end_time
    self._client: XunkemgmtClient = client
    self._universe: List[str] = universe
    self._discretionary_account_dict: Optional[Dict] = {}
    self._output_dir: str = output_dir
    self._query_discretionary_account()

  def _query_discretionary_account(self):
    with open(os.path.join(os.path.dirname(__file__), 'config/discretionary_account.json'),
              'r') as f:
      tmp = ujson.load(f)
    self._market_types = []
    self._exchanges = []
    self._owners = []
    for market_type, exchange_infos in tmp.items():
      for exchange_info, account_infos in exchange_infos.items():
        for account_info in account_infos:
          self._discretionary_account_dict[account_info['account_id']] = {
              "market_type": market_type,
              "exchange": exchange_info,
              "owner": account_info['owner']
          }
          self._market_types.append(market_type)
          self._exchanges.append(exchange_info)
          self._owners.append(account_info['owner'])
    self._market_types = list(set(self._market_types))
    self._exchanges = list(set(self._exchanges))
    self._owners = list(set(self._owners))

  def _check_subaccount_id(self, subaccount_id: int, rsp) -> List:
    res = []
    for tr_hist in rsp.histories:
      for each_transfer in tr_hist.each_transfer:
        if (each_transfer.type == cr_pb2.TransferRequestProto.TransferType.DEPOSIT_SUBACCOUNT and
            each_transfer.from_account.account_id == subaccount_id):  # sub->master
          tr_type = -1
        elif (each_transfer.type == cr_pb2.TransferRequestProto.TransferType.WITHDRAW_SUBACCOUNT and
              each_transfer.to_account.account_id == subaccount_id):  # master->sub
          tr_type = 1
        else:
          continue
        assert (each_transfer.from_account.market_type == each_transfer.to_account.market_type and
                each_transfer.from_account.exchange
                == each_transfer.to_account.exchange), "from_account and to_account are not same"
        base = each_transfer.currency
        if base not in self._universe:
          continue
        market_type = each_transfer.from_account.market_type
        exchange = each_transfer.from_account.exchange
        ts = each_transfer.create_timestamp
        qty = each_transfer.qty
        quote = MASTER_ACCOUNTING_ID[exchange]["quote"]
        product_str = f"{market_type}:{exchange}:{base}-{quote}"
        logging.info(
            f'[{subaccount_id}] {product_str} {"DEPOSIT" if tr_type == 1 else "WITHDRAW"}: {to_datetime(ts)} {qty}'
        )
        res.append(Data(to_datetime(ts), product_str, np.nan, np.nan, tr_type * qty))
    return res

  def query_daily_balance(self, account_id: int):
    res = self._client.query_balance_interval_histories(
        ais_pb2.QueryBalanceIntervalHistoriesRequestProto(
            account_ids=[account_id],
            start_timestamp=to_timestamp_int(self._start_time -
                                             datetime.timedelta(days=1)),  # since it's EOD balance
            end_timestamp=to_timestamp_int(self._end_time),
            exchanges=[self._discretionary_account_dict[account_id]["exchange"]],
            owners=[self._discretionary_account_dict[account_id]["owner"]],
        ))
    final_res = []
    for bal_hist in res.histories:
      ts = datetime.datetime.strptime(str(bal_hist.balance_info.trading_date), "%Y%m%d")
      for each_balance in bal_hist.each_balance:
        if each_balance.currency not in self._universe:
          continue
        market_type = bal_hist.balance_info.market_type
        exchange = bal_hist.balance_info.exchange
        product_str = f"{market_type}:{exchange}:{each_balance.currency}-{MASTER_ACCOUNTING_ID[exchange]['quote']}"
        final_res.append(Data(ts, product_str, np.nan, each_balance.total, np.nan))
    return final_res

  def query_real_transfer(self, account_id: int) -> List:
    exchange = self._discretionary_account_dict[account_id]["exchange"]
    master_id = MASTER_ACCOUNTING_ID[exchange]["account_id"]
    master_owner = MASTER_ACCOUNTING_ID[exchange]["owner"]
    market_type = MASTER_ACCOUNTING_ID[exchange]["market_type"]

    final_res = []
    dt = REAL_TRANSFER_START_DT
    intv = datetime.timedelta(hours=24)
    while dt <= self._end_time:
      tmp_dt = dt + intv
      res = self._client.query_transfer_histories(
          ais_pb2.QueryTransferHistoriesRequestProto(account_ids=[master_id],
                                                     start_timestamp=to_timestamp_int(dt),
                                                     end_timestamp=to_timestamp_int(tmp_dt),
                                                     market_types=[market_type],
                                                     exchanges=[exchange],
                                                     owners=[master_owner]))
      final_res.extend(self._check_subaccount_id(account_id, res))
      dt = tmp_dt
    tmp = final_res[:]
    for tr in tmp:
      ts = tr.ts
      res2 = self._client.query_balance_realtime_histories(
          ais_pb2.QueryBalanceRealtimeHistoriesRequestProto(
              account_id=account_id,
              start_timestamp=to_timestamp_int(ts - datetime.timedelta(hours=1)),
              end_timestamp=to_timestamp_int(ts + datetime.timedelta(hours=1)),
              exchange=exchange,
              owner=self._discretionary_account_dict[account_id]["owner"],
          ))
      for bal_hist in res2.histories:
        for each_balance in bal_hist.each_balance:
          if each_balance.currency not in self._universe:
            continue
          product_str = f"{market_type}:{exchange}:{each_balance.currency}-{MASTER_ACCOUNTING_ID[exchange]['quote']}"
          final_res.append(Data(ts, product_str, np.nan, each_balance.total, np.nan))
    return final_res

  def query_interval_feed(self):
    final_df = pd.DataFrame()
    start_dt = self._start_time
    end_dt = self._end_time
    assert len(set(self_product_str.split(':')[1] for self_product_str in self._product_strs)) == 1,\
      "all product_strs must be same exchange"
    for product_str in self._product_strs:
      kline_df = lac.query_interval_feed(product_str, start_dt, end_dt)
      final_df = pd.concat([final_df, kline_df])
    if final_df.empty:
      return
    final_df['ts'] = pd.to_datetime(final_df['klineTimestamp'])
    sliced_df = final_df[['ts', 'product_str', 'close']]
    sliced_df.rename(columns={'close': 'price'}, inplace=True)
    sliced_df['reserve'] = np.nan
    sliced_df['hypo_reserve'] = np.nan
    # make sliced_df to list of Data namedtuples
    final_df = sliced_df.apply(lambda x: Data(*x), axis=1).tolist()
    return final_df

  def get_corp_discretionary_pnl(self, account_id: int) -> float:
    daily_balance = self.query_daily_balance(account_id)
    real_transfer_hist = self.query_real_transfer(account_id)
    tmp = daily_balance + real_transfer_hist
    kline_universe = list(
        set([
            i.product_str
            for i in tmp
            if i.product_str.split(":")[-1].split("-")[0] not in ["USDT", "USD", "BRL", "JPY"]
        ]))
    self._product_strs = kline_universe
    kline_list = self.query_interval_feed()
    del tmp
    rtdf = pd.DataFrame(real_transfer_hist)
    if not rtdf.empty:
      rtdf = pd.DataFrame(real_transfer_hist).sort_values('ts')
      for _, group in rtdf.copy().groupby(['product_str']):
        cum_hypo_reserve = 0.0
        for idx, row in group.iterrows():
          if row['hypo_reserve'] + cum_hypo_reserve > 0:
            cum_hypo_reserve += row['hypo_reserve']
          else:
            cum_hypo_reserve = 0.0
          rtdf.loc[idx, 'hypo_reserve'] = cum_hypo_reserve
    db_df = pd.DataFrame(daily_balance)
    db_df['ts'] = db_df['ts'] + datetime.timedelta(days=1)  # Making EOD to SOD
    final_df = pd.concat([rtdf, db_df, pd.DataFrame(kline_list)],
                         axis=0).sort_values('ts').reset_index(drop=True)
    final_df.loc[final_df['ts'] == self._start_time,
                 'hypo_reserve'] = final_df.loc[final_df['ts'] == self._start_time, 'reserve']
    final_df.loc[~final_df['hypo_reserve'].isnull(),
                 'hypo_reserve'] = final_df.loc[~final_df['hypo_reserve'].isnull(), :].groupby(
                     ['product_str'])[['hypo_reserve']].cumsum()
    final_df[['price', 'hypo_reserve',
              'reserve']] = final_df.groupby(['product_str'])[['price', 'hypo_reserve',
                                                               'reserve']].ffill()
    final_df['reserve'].fillna(0.0, inplace=True)
    final_df['hypo_reserve'].fillna(0.0, inplace=True)
    final_df['corp_discretionary_pos'] = final_df['reserve'] - final_df['hypo_reserve']
    final_df[['price']] = final_df.groupby(['product_str'])[['price']].bfill()
    final_df.loc[final_df["product_str"].str.contains("USDT-USDT"), "price"] = 1.0
    final_df['price_diff'] = final_df.groupby(['product_str'])['price'].diff()
    final_df['corp_discretionary_pnl_diff'] = final_df['price_diff'] * (
        final_df['corp_discretionary_pos'])
    final_df['corp_discretionary_pnl_diff'].fillna(0.0, inplace=True)
    final_df['corp_discretionary_pnl'] = final_df.groupby(
        ['product_str'])['corp_discretionary_pnl_diff'].cumsum()
    return final_df

  def plot_corp_pnl_and_discretionary(self, account_id: int):
    df = self.get_corp_discretionary_pnl(account_id)
    df = df[~df['product_str'].str.contains("USDT-USDT")]
    os.makedirs(self._output_dir, exist_ok=True)
    exchange = self._discretionary_account_dict[account_id]['exchange']
    start_time = self._start_time.strftime('%Y%m%d')
    end_time = self._end_time.strftime('%Y%m%d')
    df.to_csv(os.path.join(self._output_dir,
                           f'{exchange}_{account_id}_{start_time}_{end_time}.csv'))
    daily_df = df.set_index('ts').groupby('product_str').resample('D').last().drop(
        'product_str', axis=1).reset_index().sort_values(['product_str', 'ts'])
    total_corp_discretionary_pnl = daily_df.groupby('ts')['corp_discretionary_pnl'].sum()
    total_corp_discretionary_pos = daily_df.groupby('ts')['corp_discretionary_pos'].sum()
    product_strs = daily_df['product_str'].unique()
    plt.rcParams['figure.figsize'] = (30, 20)
    plt.subplot(311)
    plt.plot(total_corp_discretionary_pnl.index,
             total_corp_discretionary_pnl.to_numpy(),
             'k-',
             lw=0.5,
             drawstyle='steps-post',
             label='TOTAL')
    title1 = f'corp discretionary pnl: TOTAL: {total_corp_discretionary_pnl.iloc[-1]:,.0f}, '
    for product_str in product_strs:
      tmp_df = daily_df[daily_df['product_str'] == product_str]
      plt.plot(tmp_df['ts'],
               tmp_df['corp_discretionary_pnl'],
               lw=0.5,
               drawstyle='steps-post',
               label=product_str)
      title1 += f"{product_str}: {tmp_df['corp_discretionary_pnl'].iloc[-1]:,.0f}, "
    plt.legend()
    plt.title(
        f"{title1} ({total_corp_discretionary_pnl.index.min()}-{total_corp_discretionary_pnl.index.max()})"
    )
    plt.tight_layout()
    plt.subplot(312)
    title2 = f'corp pos: '
    for product_str in product_strs:
      tmp_df = daily_df[daily_df['product_str'] == product_str]
      plt.plot(tmp_df['ts'],
               tmp_df['corp_discretionary_pos'],
               lw=0.5,
               drawstyle='steps-post',
               label=product_str)
      title2 += f"{product_str}: {tmp_df['corp_discretionary_pos'].iloc[-1]:,.0f}, "
    plt.legend()
    plt.title(
        f"{title2} ({total_corp_discretionary_pos.index.min()}-{total_corp_discretionary_pos.index.max()})"
    )
    plt.tight_layout()
    plt.subplot(313)
    for product_str in product_strs:
      tmp_df = daily_df[daily_df['product_str'] == product_str]
      plt.plot(tmp_df['ts'], tmp_df['price'], lw=0.5, drawstyle='steps-post', label=product_str)
    plt.legend()
    plt.title('price')
    plt.tight_layout()
    plt.savefig(os.path.join(self._output_dir,
                             f'{exchange}_{account_id}_{start_time}_{end_time}.png'),
                dpi=300)
    plt.close()

  def plot_corp_pnl_and_discretionary_all(self):
    for account_id in self._discretionary_account_dict:
      try:
        self.plot_corp_pnl_and_discretionary(account_id)
      except Exception:
        logging.error(f"failed to plot {account_id}: {traceback.format_exc()}")


def main(_):
  start_dt = datetime.datetime.strptime(FLAGS.start_time, '%Y%m%d')
  end_dt = datetime.datetime.strptime(FLAGS.end_time, '%Y%m%d')
  account_id = "ALL" if FLAGS.account_id == "ALL" else int(FLAGS.account_id)
  output_dir = os.path.join(os.path.expandvars('$COIN_REPO'), FLAGS.output_dir)
  universe = FLAGS.universe.split(',') if FLAGS.universe else None
  with XunkemgmtClient() as client:
    discretionary = Discretionary(start_dt, end_dt, client, universe, output_dir)
    if account_id != "ALL":
      discretionary.plot_corp_pnl_and_discretionary(account_id)
    else:
      discretionary.plot_corp_pnl_and_discretionary_all()


if __name__ == "__main__":
  flags.DEFINE_string('start_time', '20230101', '%Y%m%d')
  flags.DEFINE_string('end_time', '20230331', '%Y%m%d')
  flags.DEFINE_string('account_id', 'ALL', 'account_id')
  flags.DEFINE_string('output_dir', 'sim_result/accounting', 'output dir')
  flags.DEFINE_string('universe', '', 'discretionary trading universe')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
